{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e0b9a6ba-ca22-4529-a0ba-2c963ddf668a",
   "metadata": {},
   "source": [
    "# Customizing Lennard-Jones Mixing\n",
    "\n",
    "OpenMM's standard [NonbondedForce](https://docs.openmm.org/latest/api-python/generated/openmm.openmm.NonbondedForce.html) class for computing Lennard-Jones and electrostatic interactions implements the Lorentz-Berthelot mixing rules for Lennard-Jones parameters.  As described in the [user guide](https://docs.openmm.org/latest/userguide/theory/02_standard_forces.html#lennard-jones-interaction), after parameters $\\sigma_i$ and $\\epsilon_i$ have been specified for *each* particle $i$ in an OpenMM System, the actual Lennard-Jones parameters used to compute the interaction between two particles $i$ and $j$ are\n",
    "$$\\begin{align}\\sigma_{ij}&=\\frac{\\sigma_i+\\sigma_j}{2}\\\\\\epsilon_{ij}&=\\sqrt{\\epsilon_i\\epsilon_j}.\\end{align}$$\n",
    "You may want to use different mixing rules for Lennard-Jones interactions in a system, or fully customize the Lennard-Jones parameters employed for interactions between pairs of particles in different sets.  NonbondedForce supports [adding exceptions](https://docs.openmm.org/latest/api-python/generated/openmm.openmm.NonbondedForce.html#openmm.openmm.NonbondedForce.addException) to specific pairs of particles, but exceptions should not be used to implement custom mixing rules, since exceptions are computed like *bonded* interactions that do not respect cutoffs.\n",
    "\n",
    "Instead, you can use [CustomNonbondedForce](https://docs.openmm.org/latest/api-python/generated/openmm.openmm.CustomNonbondedForce.html) to implement the Lennard-Jones interactions.  To handle the typical most general case, in which each particle is assigned one of several particle types, and every pair of particle types may have its own set of parameters, you can use a [Discrete2DFunction](https://docs.openmm.org/latest/api-python/generated/openmm.openmm.Discrete2DFunction.html) to store tables of parameter values.  To demonstrate this, we first create a [System](https://docs.openmm.org/latest/api-python/generated/openmm.openmm.System.html) with some particles."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "fecb6cec-beaa-4273-8828-aad9c243b712",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import openmm\n",
    "\n",
    "system = openmm.System()\n",
    "for i in range(300):\n",
    "    system.addParticle(1.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca5b90a5-0edc-4a92-bb89-ff1fa244131c",
   "metadata": {},
   "source": [
    "Now we create a CustomNonbondedForce representing a Lennard-Jones interaction, and add it to the System:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0be9403f-6ce1-4582-8e64-bcda4b6301aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "lj = openmm.CustomNonbondedForce(\"4*epsilon(type1, type2)*((sigma(type1, type2)/r)^12 - (sigma(type1, type2)/r)^6)\")\n",
    "system.addForce(lj);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9ae9b2c-b186-4f6e-93f8-a74ec4a473f1",
   "metadata": {},
   "source": [
    "Note that OpenMM has no built-in notion of \"particle types\", but by adding a per-particle parameter called `type` using [CustomNonbondedForce.addPerParticleParameter()](https://docs.openmm.org/latest/api-python/generated/openmm.openmm.CustomNonbondedForce.html#openmm.openmm.CustomNonbondedForce.addPerParticleParameter), the names `type1` and `type2` in the custom expression will be automatically associated with the values of the `type` parameter for pairs of interacting particles."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1744000c-d9bd-4ee9-b82c-fedec1749d4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "lj.addPerParticleParameter(\"type\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6534330-c8c7-4390-81f9-7f44ee3171ae",
   "metadata": {},
   "source": [
    "To assign the actual `type` values for each particle, we can use [CustomNonbondedForce.addParticle()](https://docs.openmm.org/latest/api-python/generated/openmm.openmm.CustomNonbondedForce.html#openmm.openmm.CustomNonbondedForce.addParticle).  To make up some values for this example, we'll set the type of the first 150 particles to 0, the next 100 to 1, and the last 50 to 2."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9ea1c4ac-2efd-4307-a8d2-1a9324e1045f",
   "metadata": {},
   "outputs": [],
   "source": [
    "particle_types = [0] * 150 + [1] * 100 + [2] * 50\n",
    "for particle_type in particle_types:\n",
    "    lj.addParticle([particle_type])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a44baad-9ee4-4873-872b-edbf4f740cdb",
   "metadata": {},
   "source": [
    "Using [CustomNonbondedForce.addTabulatedFunction()](https://docs.openmm.org/latest/api-python/generated/openmm.openmm.CustomNonbondedForce.html#openmm.openmm.CustomNonbondedForce.addTabulatedFunction), we can then add tabulated functions for `sigma` and `epsilon` so that values for these parameters can be looked up by pairs of particle type indices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "30f67287-d539-45c7-a086-44e7dc5ea4ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "lj.addTabulatedFunction(\"sigma\", openmm.Discrete2DFunction(3, 3, np.array([\n",
    "    [1.0, 1.3, 1.4],\n",
    "    [1.3, 1.1, 1.5],\n",
    "    [1.4, 1.5, 1.2]\n",
    "]).flatten(order=\"F\")))\n",
    "lj.addTabulatedFunction(\"epsilon\", openmm.Discrete2DFunction(3, 3, np.array([\n",
    "    [1.0, 1.5, 2.0],\n",
    "    [1.5, 2.0, 2.5],\n",
    "    [2.0, 2.5, 3.0]\n",
    "]).flatten(order=\"F\")));"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2d9a663-f08c-4d86-b23b-99cb65a0f9be",
   "metadata": {},
   "source": [
    "Note that:\n",
    "* The matrices of parameters are symmetric, so that, *e.g.*, `sigma(0, 1)` and `sigma(1, 0)` are equal.  When CustomNonbondedForce's energy expression is evaluated, pairs of particles may appear in any order, and there is no guarantee that they will be ordered consistently.  Thus, the energy expression *must* be symmetric with respect to exchange of particles to obtain sensible results, and so in this example, these matrices must be symmetric.\n",
    "* [Discrete2DFunction](https://docs.openmm.org/latest/api-python/generated/openmm.openmm.Discrete2DFunction.html) takes a flat array of values, in *column-major* (*i.e.*, \"Fortran\") order.  Here, therefore, we use 2D NumPy arrays, and flatten them with `order=\"F\"`.  Since the arrays are symmetric in this example, this is not strictly needed, but it is a good practice to avoid surprises working with multidimensional tabulated functions in OpenMM.\n",
    "* This example assumes the use of dimensionless units.  Quantities with units passed to OpenMM's custom force and tabulated function classes will be converted to OpenMM's [default units](https://docs.openmm.org/latest/userguide/theory/01_introduction.html#units) of nm, ps, amu, kJ/mol, *etc.*  Results when using units should thus be valid as long as all custom expressions given to OpenMM are dimensionally consistent.  *It is the responsibility of the user* to ensure that they are dimensionally consistent, as OpenMM does not verify this.\n",
    "\n",
    "At this point, we could set any other options we desire on the CustomNonbondedForce, and set up other force field terms in the System as usual.  CustomNonbondedForce supports several functions that the standard NonbondedForce supports, including [setCutoffDistance()](https://docs.openmm.org/latest/api-python/generated/openmm.openmm.CustomNonbondedForce.html#openmm.openmm.CustomNonbondedForce.setCutoffDistance), [setNonbondedMethod()](https://docs.openmm.org/latest/api-python/generated/openmm.openmm.CustomNonbondedForce.html#openmm.openmm.CustomNonbondedForce.setNonbondedMethod), [setSwitchingDistance()](https://docs.openmm.org/latest/api-python/generated/openmm.openmm.CustomNonbondedForce.html#openmm.openmm.CustomNonbondedForce.setSwitchingDistance), [setUseLongRangeCorrection()](https://docs.openmm.org/latest/api-python/generated/openmm.openmm.CustomNonbondedForce.html#openmm.openmm.CustomNonbondedForce.setUseLongRangeCorrection), and [setUseSwitchingFunction()](https://docs.openmm.org/latest/api-python/generated/openmm.openmm.CustomNonbondedForce.html#openmm.openmm.CustomNonbondedForce.setUseSwitchingFunction).  Keep in mind that if your system includes both Lennard-Jones interactions with custom mixing rules *and* electrostatic interactions, you should use NonbondedForce for the electrostatics, setting all sigma and epsilon values to zero in the NonbondedForce, and then add a CustomNonbondedForce as described here."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3097b65d-98de-4972-90d4-56203cdfd353",
   "metadata": {},
   "source": [
    "## Custom Expressions\n",
    "\n",
    "You can populate the `sigma` and `epsilon` tables with arbitrary values, or even use some other mixing rule to calculate them.  If, instead of using tabulated values for all particle type pairs, you want to have OpenMM compute values according to some mixing rule from `sigma` and `epsilon` values set for each *particle*, you can include your rule in the CustomNonbondedForce energy expression, and add per-particle parameters for these values instead of particle types, *e.g.*:\n",
    "```python\n",
    "lj = openmm.CustomNonbondedForce(\n",
    "    \"4*epsilon*((sigma/r)^12 - (sigma/r)^6);\"\n",
    "\n",
    "    # Expressions can depend on sigma1, sigma2, epsilon1, and epsilon2\n",
    "    # and must be symmetric with respect to exchange of 1 and 2\n",
    "    \"sigma = ...; epsilon = ...\"\n",
    ")\n",
    "lj.addPerParticleParameter(\"sigma\")\n",
    "lj.addPerParticleParameter(\"epsilon\")\n",
    "\n",
    "for ...:\n",
    "    # Specify sigma and epsilon for each particle\n",
    "    lj.addParticle([..., ...])\n",
    "```\n",
    "Since you can use any energy expression in CustomNonbondedForce, both this approach and the tabular approach above can even be extended to pairwise interactions with arbitrary functional forms, not just Lennard-Jones interactions."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65fa041e-7543-4611-b062-65ccae417437",
   "metadata": {},
   "source": [
    "## Exceptions and Exclusions\n",
    "\n",
    "Unlike NonbondedForce, CustomNonbondedForce does not support \"exceptions\" that permit you to *modify* interaction parameters for particular pairs of particles, only \"exclusions\" that allow you to omit pairs entirely.  If you only need to zero the Lennard-Jones parameters for certain pairs, you can use [CustomNonbondedForce.addExclusion()](https://docs.openmm.org/latest/api-python/generated/openmm.openmm.CustomNonbondedForce.html#openmm.openmm.CustomNonbondedForce.addExclusion) or [CustomNonbondedForce.createExclusionsFromBonds()](https://docs.openmm.org/latest/api-python/generated/openmm.openmm.CustomNonbondedForce.html#openmm.openmm.CustomNonbondedForce.createExclusionsFromBonds).  If you also want to include modified interactions, you can exclude them from the CustomNonbondedForce and add a [CustomBondForce](https://docs.openmm.org/latest/api-python/generated/openmm.openmm.CustomBondForce.html) implementing the Lennard-Jones functional form for the pairs of interest.  For instance, this could look like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "1812cb4b-d0c9-4108-adb2-7a347a0290a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "lj_exceptions = openmm.CustomBondForce(\"4*epsilon*((sigma/r)^12 - (sigma/r)^6)\")\n",
    "lj_exceptions.addPerBondParameter(\"sigma\")\n",
    "lj_exceptions.addPerBondParameter(\"epsilon\")\n",
    "system.addForce(lj_exceptions)\n",
    "\n",
    "def add_exception(i, j, epsilon, sigma):\n",
    "    lj.addExclusion(i, j)\n",
    "\n",
    "    # For LJ interactions, if either sigma or epsilon is zero, the\n",
    "    # exception force will evaluate to zero, so no bond is needed.\n",
    "    if sigma and epsilon:\n",
    "        lj_exceptions.addBond(i, j, [sigma, epsilon])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.14.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}